package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"log"
	"math/big"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"time"
)

func generateSelfSignedCert(ip string, port string, org string, ou string) (string, string, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", err
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(365 * 24 * time.Hour) // 一年有效期

	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return "", "", err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization:       []string{org},
			OrganizationalUnit: []string{ou},
			CommonName:         ip + ":" + port,
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return "", "", err
	}

	certOut, err := os.Create("server.crt")
	if err != nil {
		return "", "", err
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	certOut.Close()

	keyOut, err := os.Create("server.key")
	if err != nil {
		return "", "", err
	}
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", err
	}
	pem.Encode(keyOut, &pem.Block{Type: "PRIVATE KEY", Bytes: privBytes})
	keyOut.Close()

	return "server.crt", "server.key", nil
}

func main() {
	// 解析命令行参数
	ip := flag.String("ip", "0.0.0.0", "IP地址")
	port := flag.String("port", "1443", "端口号")
	targetURL := flag.String("target", "http://10.2.2.1:8089/", "目标URL")
	org := flag.String("org", "go-proxy-https", "证书组织名称")
	ou := flag.String("ou", "太杀猪刀", "证书组织单位名称")
	flag.Parse()

	// 检查目标URL是否为空
	if *targetURL == "" {
		log.SetFlags(log.LstdFlags | log.Lshortfile)
		log.SetPrefix("错误: ")
		log.Fatalf("目标URL不能为空")
	}

	// 自动生成证书和私钥
	certFile, keyFile, err := generateSelfSignedCert(*ip, *port, *org, *ou)
	if err != nil {
		log.SetFlags(log.LstdFlags | log.Lshortfile)
		log.SetPrefix("错误: ")
		log.Fatalf("生成证书时出错: %v", err)
	}

	// 加载证书和私钥
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.SetFlags(log.LstdFlags | log.Lshortfile)
		log.SetPrefix("错误: ")
		log.Fatalf("加载证书和私钥时出错: %v", err)
	}

	// 创建TLS配置
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	// 创建TLS监听器
	listener, err := tls.Listen("tcp", *ip+":"+*port, tlsConfig)
	if err != nil {
		log.SetFlags(log.LstdFlags | log.Lshortfile)
		log.SetPrefix("错误: ")
		log.Fatalf("监听 %s:%s 时出错: %v", *ip, *port, err)
	}
	defer listener.Close()

	// 创建HTTP服务
	mux := http.NewServeMux()

	mux.HandleFunc("/error", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("发生错误"))
	})

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// 将请求转发到HTTP服务
		target, _ := url.Parse(*targetURL)
		proxy := httputil.NewSingleHostReverseProxy(target)

		// 将请求转发到目标 URL
		proxy.ServeHTTP(w, r)
	})

	// 启动HTTPS服务
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.SetPrefix("信息: ")
	log.Printf("正在以 HTTPS 服务运行在 %s:%s", *ip, *port)
	if err := http.Serve(listener, mux); err != nil {
		log.SetFlags(log.LstdFlags | log.Lshortfile)
		log.SetPrefix("错误: ")
		log.Fatalf("启动 HTTPS 服务时出错: %v", err)
	}
}
